use std::{fs::{File, OpenOptions}, io::Write};

use super::{concat, SemanticAnalyzer, HashMap, hashmap, TokenType, AST, GLOBAL, MAIN, INDEX_OUT_OF_BOUNDS, Instruction};

#[derive(Debug, Clone, Default)]
pub struct CodeGenerator {
    _root: Box<AST>,
    _symbol_table: Box<HashMap<String, HashMap<String, (i32, i32)>>>,
}

impl CodeGenerator {
    pub fn new(root: &AST, symbol_table: &HashMap<String, HashMap<String, (i32, i32)>>) -> Self {
        Self {
            _root: Box::from(root.clone()),
            _symbol_table: Box::from(symbol_table.clone()),
        }
    }
    /// Generate Code
    pub fn generate_code(&mut self) -> Vec<(Instruction, String)> {
        let mut code_map = self.create_code_map();
        let (mut code_list, func_jmp_map) = Self::merge_code_map(&code_map);

        self.translate_call(&mut code_list, &func_jmp_map);

        code_list
    }
    /// 代码装载
    fn create_code_map(&mut self) -> HashMap<String, Vec<(Instruction, String)>> {
        let mut code_map = hashmap!(String::from(GLOBAL) => self.gen_global_code());

        for decl_ptr in &*self._root.children {
            if decl_ptr.token_type == TokenType::FuncDecl {
                let cur_func_name = decl_ptr
                                    .children
                                    .get(1)
                                    .expect(INDEX_OUT_OF_BOUNDS)
                                    .token_str.as_ref();
                let mut code_list = self.gen_stmt_list_code(
                    decl_ptr
                        .children
                        .get(3)
                        .expect(INDEX_OUT_OF_BOUNDS)
                        .children
                        .get(1)
                        .expect(INDEX_OUT_OF_BOUNDS),
                        cur_func_name
                );
                if cur_func_name != MAIN {
                    code_list.push((Instruction::RET, String::default()))
                }
                code_map.insert(cur_func_name.clone(), code_list);
            }
        }
        code_map
    }
    fn gen_number_code(&self, root: &AST) -> Vec<(Instruction, String)> {
        vec![(Instruction::LDC, (*root.token_str).clone())]
    }
    fn gen_compound_stmt_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        self.gen_stmt_list_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name)
    }
    fn gen_stmt_list_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let mut code_list = Vec::default();
        for child in &*root.children {
            let mut stmt_code_list = self.gen_stmt_code(child, cur_func_name);
            code_list.append(&mut stmt_code_list);
        }
        code_list
    }
    fn gen_stmt_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        match root.token_type {
            TokenType::Expr => self.gen_expr_code(root, cur_func_name),
            TokenType::CompoundStmt => self.gen_compound_stmt_code(root, cur_func_name),
            TokenType::IfStmt => self.gen_if_stmt_code(root, cur_func_name),
            TokenType::WhileStmt => self.gen_while_stmt_code(root, cur_func_name),
            TokenType::ReturnStmt => self.gen_return_stmt_code(root, cur_func_name),
            _ => panic!("Invalid TokenType"),
        }
    }
    fn gen_if_stmt_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let mut code_list = self.gen_expr_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
        let mut if_code_list = self.gen_stmt_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);

        if root.children.len() == 2 {
            code_list.push((Instruction::JZ, (if_code_list.len() + 1).to_string()));
            code_list.append(&mut if_code_list);
        } else {
            let mut else_code_list =
                self.gen_stmt_code(root.children.get(2).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);

            if_code_list.push((Instruction::JMP, (else_code_list.len() + 1).to_string()));
            code_list.push((Instruction::JZ, (if_code_list.len() + 1).to_string()));
            code_list.append(&mut if_code_list);
            code_list.append(&mut else_code_list);
        }
        code_list
    }
    fn gen_while_stmt_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let mut code_list = self.gen_expr_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
        let mut stmt_code_list =
            self.gen_stmt_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);

        code_list.push((Instruction::JZ, (stmt_code_list.len() + 2).to_string()));
        code_list.append(&mut stmt_code_list);
        code_list.push((Instruction::JMP, concat("-", code_list.len())));
        code_list
    }
    fn gen_return_stmt_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        if root.children.len() == 0 {
            Vec::default()
        } else {
            self.gen_expr_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name)
        }
    }
    fn gen_expr_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        if root.children.len() == 1 {
            self.gen_simple_expr_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name)
        } else {
            let mut code_list =
                self.gen_expr_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
            let mut assign_code_list =
                self.gen_assign_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
            code_list.append(&mut assign_code_list);
            code_list
        }
    }
    fn gen_var_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let mut code_list = Vec::default();

        let symbol_table = &self._symbol_table;
        // let token_str = &*root.children.get(0).expect(INDEX_OUT_OF_BOUNDS).token_str;
        // let child_of_child = child
        //     .get(cur_func_name)
        //     .expect(INDEX_OUT_OF_BOUNDS)
        //     .get(token_str);
        let child_of_child = symbol_table[cur_func_name].get(&*root.children[0].token_str);
        if child_of_child.is_some() {
            code_list.push((Instruction::LDC, child_of_child.unwrap().0.to_string()));
            code_list.push((Instruction::LD, String::default()));
        } else {
            code_list.push((
                Instruction::LDC,
                symbol_table[GLOBAL][&*root.children[0].token_str].0.to_string(),
            ));
            code_list.push((Instruction::ALD, String::default()));
        }

        if root.children.len() == 2 {
            let mut expr_code_list =
                self.gen_expr_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
            code_list.push((Instruction::PUSH, Default::default()));
            code_list.append(&mut expr_code_list);
            code_list.push((Instruction::ADD, Default::default()));
            code_list.push((Instruction::POP, Default::default()));
            code_list.push((Instruction::ALD, Default::default()));
        }
        code_list
    }
    fn gen_simple_expr_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        if root.children.len() == 1 {
            self.gen_add_expr_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name)
        } else {
            let mut code_list =
                self.gen_add_expr_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
            let mut mid_code_list =
                self.gen_rel_op_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS));
            let mut right_code_list =
                self.gen_add_expr_code(root.children.get(2).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);

            code_list.push((Instruction::PUSH, Default::default()));
            code_list.append(&mut right_code_list);
            code_list.append(&mut mid_code_list);
            code_list.push((Instruction::POP, Default::default()));

            code_list
        }
    }
    fn gen_rel_op_code(&self, root: &AST) -> Vec<(Instruction, String)> {
        match root.token_type {
            TokenType::Less => vec![(Instruction::LT, Default::default())],
            TokenType::LessEq => vec![(Instruction::LE, Default::default())],
            TokenType::Greater => vec![(Instruction::GT, Default::default())],
            TokenType::GreaterEq => vec![(Instruction::GE, Default::default())],
            TokenType::Eq => vec![(Instruction::EQ, Default::default())],
            TokenType::NotEq => vec![(Instruction::NE, Default::default())],
            _ => panic!("Invalid TokenType"),
        }
    }
    fn gen_add_expr_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let mut code_list = self.gen_term_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
        let mut i = 1;
        while i < root.children.len() {
            let mut mid_code_list =
                self.gen_add_op_code(root.children.get(i).expect(INDEX_OUT_OF_BOUNDS));
            let mut right_code_list =
                self.gen_term_code(root.children.get(i + 1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);

            code_list.push((Instruction::PUSH, Default::default()));
            code_list.append(&mut right_code_list);
            code_list.append(&mut mid_code_list);
            code_list.push((Instruction::POP, Default::default()));

            i += 2;
        }
        code_list
    }
    fn gen_add_op_code(&self, root: &AST) -> Vec<(Instruction, String)> {
        match root.token_type {
            TokenType::Plus => vec![(Instruction::ADD, Default::default())],
            _ => vec![(Instruction::SUB, Default::default())],
        }
    }
    fn gen_term_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let mut code_list = self.gen_factor_code(root.children.get(0).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
        let mut i = 1;
        while i < root.children.len() {
            let mut mid_code_list =
                self.gen_mul_op_code(root.children.get(i).expect(INDEX_OUT_OF_BOUNDS));
            let mut right_code_list =
                self.gen_factor_code(root.children.get(i + 1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);

            code_list.push((Instruction::PUSH, Default::default()));
            code_list.append(&mut right_code_list);
            code_list.append(&mut mid_code_list);
            code_list.push((Instruction::POP, Default::default()));

            i += 2;
        }
        code_list
    }
    fn gen_mul_op_code(&self, root: &AST) -> Vec<(Instruction, String)> {
        match root.token_type {
            TokenType::Multiply => vec![(Instruction::MUL, Default::default())],
            _ => vec![(Instruction::DIV, Default::default())],
        }
    }
    fn gen_factor_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        match root.token_type {
            TokenType::Expr => self.gen_expr_code(root, cur_func_name),
            TokenType::Number => self.gen_number_code(root),
            TokenType::Call => self.gen_call_code(root, cur_func_name),
            TokenType::Var => self.gen_var_code(root, cur_func_name),
            _ => panic!("Invalid TokenType"),
        }
    }
    fn gen_call_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let token_str = &*root.children.get(0).expect(INDEX_OUT_OF_BOUNDS).token_str;
        //内置函数
        if token_str == "input" {
            //input
            return vec![(Instruction::IN, String::default())];
        } else if token_str == "output" {
            //output
            let mut code_list = self.gen_expr_code(
                root.children
                    .get(1)
                    .expect(INDEX_OUT_OF_BOUNDS)
                    .children
                    .get(0)
                    .expect(INDEX_OUT_OF_BOUNDS),
                    cur_func_name
            );
            code_list.push((Instruction::OUT, String::default()));
            return code_list;
        }
        // 从符号表中分离出所有的局部变量信息
        let mut code_list = Vec::default();
        let symbol_table = self
            ._symbol_table
            .get(&*root.children.get(0).expect(INDEX_OUT_OF_BOUNDS).token_str)
            .expect(INDEX_OUT_OF_BOUNDS);
        let size = symbol_table.len();

        let mut pair_list = Vec::default();
        for _ in 0..size {
            pair_list.push((String::default(), (i32::default(), i32::default())));
        }
        for (first, &(second_one, second_two)) in symbol_table {
            let index = (size as i32) - second_one - 1;
            if index < 0 {continue;}
            pair_list[index as usize] = (first.clone(), (second_one, second_two));
        }
        let top_idx = size
            - (if root.children.len() == 2 {
                root.children
                    .get(1)
                    .expect(INDEX_OUT_OF_BOUNDS)
                    .children
                    .len()
            } else {
                0
            });
        // 局部变量压栈
        for i in 0..top_idx {
            let tmp = pair_list.get(i).expect(INDEX_OUT_OF_BOUNDS).1.1;
            if tmp > 0 {
                for _ in 0..tmp {
                    code_list.push((Instruction::PUSH, Default::default()));
                }
                code_list.push((Instruction::ADDR, Default::default()));
                code_list.push((Instruction::PUSH, Default::default()));
            } else {
                code_list.push((Instruction::PUSH, Default::default()));
            }
        }
        // 实参压栈
        if root.children.len() == 2 {
            let mut arg_list_code_list =
                self.gen_arg_list_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
            code_list.append(&mut arg_list_code_list);
        }
        // BP寄存器准备
        // IP寄存器准备与函数调用
        code_list.push((Instruction::CALL, *root.children[0].token_str.clone()));
        // 实参与局部变量退栈
        for child in symbol_table {
            code_list.push((Instruction::POP, String::default()));

            for _ in 0..child.1.1 {
                code_list.push((Instruction::POP, String::default()));
            }
        }
        code_list
    }
    fn gen_arg_list_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let mut code_list = Vec::default();
        for i in (0..root.children.len()).rev() {
            let mut expr_code_list =
                self.gen_expr_code(root.children.get(i).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
            code_list.append(&mut expr_code_list);
            code_list.push((Instruction::PUSH, Default::default()));
        }
        code_list
    }
    fn gen_assign_code(&self, root: &AST, cur_func_name: &String) -> Vec<(Instruction, String)> {
        let mut code_list = vec![(Instruction::PUSH, Default::default())];

        let child = &self._symbol_table;
        let token_str = &*root.children.get(0).expect(INDEX_OUT_OF_BOUNDS).token_str;
        let child_of_child = child
            .get(cur_func_name)
            .expect(INDEX_OUT_OF_BOUNDS)
            .get(token_str);
        if child_of_child.is_some() {
            code_list.push((Instruction::LDC, child_of_child.unwrap().0.to_string()));
            if root.children.len() == 1 {
                code_list.push((Instruction::ST, Default::default()));
            } else {
                let mut expr_code_list =
                    self.gen_expr_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
                // Get the (start) pointer (Is already an absolute address)
                code_list.push((Instruction::LD, Default::default()));
                code_list.push((Instruction::PUSH, Default::default()));
                code_list.append(&mut expr_code_list);

                // Pointer[Index] (Pointer + Index)
                code_list.push((Instruction::ADD, Default::default()));
                code_list.push((Instruction::POP, Default::default()));

                // Save by absolute address
                code_list.push((Instruction::AST, Default::default()));
            }
        } else {
            code_list.push((
                Instruction::LDC,
                child
                    .get(GLOBAL)
                    .expect(INDEX_OUT_OF_BOUNDS)
                    .get(token_str)
                    .unwrap_or(&(i32::default(), i32::default()))
                    .0.to_string(),
            ));
            if root.children.len() == 1 {
                code_list.push((Instruction::AST, Default::default()));
            } else {
                let mut expr_code_list =
                    self.gen_expr_code(root.children.get(1).expect(INDEX_OUT_OF_BOUNDS), cur_func_name);
                code_list.push((Instruction::ALD, Default::default()));
                code_list.push((Instruction::PUSH, Default::default()));
                code_list.append(&mut expr_code_list);
                code_list.push((Instruction::ADD, Default::default()));
                code_list.push((Instruction::POP, Default::default()));
                code_list.push((Instruction::AST, Default::default()));
            }
        }

        code_list.push((Instruction::POP, Default::default()));
        code_list
    }
    /// 链接器1
    fn create_func_jmp_map(
        &self,
        code_map: &HashMap<String, Vec<(Instruction, String)>>,
    ) -> HashMap<String, i32> {
        let mut func_jmp_map = hashmap!(String::from(GLOBAL) => 0);
        let mut jmp_num = code_map.get(GLOBAL).expect("no global").len() as i32;

        for (first, second) in code_map {
            if first != GLOBAL && first != MAIN {
                func_jmp_map.insert(first.clone(), jmp_num);
                jmp_num += second.len() as i32;
            }
        }

        func_jmp_map.insert(String::from(MAIN), jmp_num);

        func_jmp_map
    }
    fn translate_call_helper(&self, code_list: &mut Vec<(Instruction, String)>, ip: &mut i32, func_jmp_map: &HashMap<String, i32>) {
        for (code_enum, code_val_str) in code_list {
            if *code_enum == Instruction::CALL {
                *code_val_str = (*func_jmp_map.get(code_val_str).unwrap() - *ip).to_string();
            }
            *ip += 1;
        }
    }
    fn merge_code_map(code_map: &HashMap<String, Vec<(Instruction, String)>>) -> (Vec<(Instruction, String)>, HashMap<String, i32>) {
        let mut code_list = Vec::<(Instruction, String)>::default();
        let mut func_jmp_map = hashmap!(GLOBAL.to_string() => 0);
        let mut code_map_global = code_map.get(GLOBAL).unwrap().clone();
        let mut jmp_num = code_map_global.len() as i32;
        code_list.append(&mut code_map_global);

        for (func_name, sub_code_list) in code_map {
            if func_name != GLOBAL && func_name != MAIN {
                code_list.append(&mut sub_code_list.clone());
                func_jmp_map.insert(func_name.clone(), jmp_num);
                jmp_num += sub_code_list.len() as i32;
            }
        }

        let mut code_map_main = code_map.get(MAIN).unwrap().clone();
        code_list.append(&mut code_map_main);

        func_jmp_map.insert(MAIN.to_string(), jmp_num);
        (code_list, func_jmp_map)
    }
    /// 链接器2
    fn translate_call(&self, code_list: &mut Vec<(Instruction, String)>, func_jmp_map: &HashMap<String, i32>) {
        for ip in 0..code_list.len() {
            if code_list[ip].0 == Instruction::CALL {
                code_list[ip].1 = (func_jmp_map[&code_list[ip].1] - (ip as i32)).to_string();
            }
        }
    }
    /// 全局变量压栈
    fn gen_global_variable_code(&self) -> Vec<(Instruction, String)> {
        let mut code_list = Vec::default();
        let symbol_table = self._symbol_table.get(GLOBAL).expect("no global");
        for (_, &(second_one, second_two)) in symbol_table {
            if second_two != 0 {
                code_list.push((Instruction::LDC, (second_one + 1).to_string()));
            }
            code_list.push((Instruction::PUSH, Default::default()));
            for _ in 0..second_two {
                code_list.push((Instruction::PUSH, Default::default()));
            }
        }
        code_list
    }
    // main函数
    fn gen_main_prepare_code(&self) -> Vec<(Instruction, String)> {
        let mut code_list = Vec::default();
        let main = self._symbol_table.get(MAIN).expect("no main");
        
        let mut pair_list = Vec::<(String, (i32, i32))>::new();
        for _ in 0..main.len() {
            pair_list.push((String::default(), (i32::default(), i32::default())));
        }
        for (first, &(second_one, second_two)) in main {
            // let pair_list_clone = pair_list.clone();
            // pair_list_clone.len() - second_one - 1;
            //此处不可能越界
            let index = (pair_list.len() as i32) - second_one - 1;
            // if index < 0 {continue;}
            pair_list[index as usize] = (first.clone(), (second_one, second_two));
        }
        for (first, (second_one, second_two)) in pair_list {
            if second_two != 0 {
                for _ in 0..second_two {
                    code_list.push((Instruction::PUSH, Default::default()));
                }
                code_list.push((Instruction::ADDR, second_two.to_string()));
                code_list.push((Instruction::PUSH, Default::default()));
            } else {
                code_list.push((Instruction::PUSH, Default::default()));
            }
        }
        code_list.push((Instruction::CALL, String::from(MAIN)));
        code_list
    }

    fn gen_global_code(&self) -> Vec<(Instruction, String)>{
        let mut code_list = self.gen_global_variable_code();
        let mut main_prepare_code_list = self.gen_main_prepare_code();

        code_list.append(&mut main_prepare_code_list);
        code_list
    }

}
